package com.test.html;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import okhttp3.*;
import okio.BufferedSink;
import okio.Okio;
import okio.Sink;
import org.apache.commons.lang3.StringUtils;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.util.Scanner;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 下载midjourney图片
 */
public class HtmlParser3 {
    private static AtomicInteger handled = new AtomicInteger(0);

    private static LocalDate now = LocalDate.now();
    private static String date = now.format(DateTimeFormatter.ofPattern("yy-MM"));
    private static String baseDir = "I:\\ai\\midjourney\\" + date;
    private static String logFileName = baseDir + "_" + "jsonData.txt";

    // 信号量：6，公平锁
    private static Semaphore semaphore = new Semaphore(6, true);

    public static void main(String[] args) throws IOException {
        System.out.println(baseDir);
        Path path = Paths.get(baseDir);
        Files.createDirectories(path);

        /* 下载单个图片
        download("https://cdn.midjourney.com/1b9f6d82-a5dd-4600-98d8-3a54c637a7bd/0_1.png");
        try {
            Thread.sleep(4000L);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }*/

        // 从midjourney下载图片
        /*String url = // "https://www.midjourney.com/showcase/top/"; //
                "https://www.midjourney.com/showcase/recent/";
        String doc = OkHttpUtilsLocal.builder()
                .url(url)
                .addHeader("authority","cdn.midjourney.com")
                .addHeader("accept","text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng;q=0.8,application/signed-exchange;v=b3;q=0.7")
//                .addHeader("accept-language","zh-CN,zh;q=0.9")
                .addHeader("sec-ch-ua","\"Not.A/Brand\";v=\"8\", \"Chromium\";v=\"114\", \"Google Chrome\";v=\"114\"")
                .addHeader("sec-ch-ua-mobile","?0")
                .addHeader("sec-ch-ua-platform","\"Windows\"")
                .addHeader("sec-fetch-dest","document")
                .addHeader("sec-fetch-mode","navigate")
                .addHeader("sec-fetch-site","none")
                .addHeader("sec-fetch-user","?0")
                .addHeader("upgrade-insecure-requests","0")
                .addHeader("user-agent","Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36")
//                .addHeader("cookie","_ga=GA1.1.862279448.1680306197; _ga_Q0DQ5L7K0D=GS1.1.1685918653.42.0.1685918653.0.0.0")
                .get()
                .sync();
        System.out.println(doc);
        String json = doc.substring(doc.indexOf("<script id=\"__NEXT_DATA__\" type=\"application/json\">") + 51);
        json  = json.substring(0, doc.lastIndexOf("</script>")); */
        Path jsonPath = Paths.get("I:\\ai\\midjourney\\recent.txt");
        Scanner scanner = new Scanner(jsonPath);
        StringBuilder json = new StringBuilder("");
        while (scanner.hasNextLine()) {
            json.append(scanner.nextLine());
        }
        System.out.println(json.toString());
        saveToFile(logFileName, now.format(DateTimeFormatter.ofPattern("yy-MM-dd")));
        JSONObject data = JSONObject.parseObject(json.toString());
        JSONObject props = data.getJSONObject("props");
        JSONObject pageProps = props == null ? data.getJSONObject("pageProps") : props.getJSONObject("pageProps");
        JSONArray jobs = pageProps.getJSONArray("jobs");
        int size = jobs.size();
        JSONObject job;
        String imageUrl;
        for (int i = 0; i < size; i++) {
            job = jobs.getJSONObject(i);
            imageUrl = job.getJSONObject("event").getString("seedImageURL");
            if (imageUrl == null) {
                imageUrl = job.getJSONArray("image_paths").getString(0);
            }
            String fileName = getFileNameFromUrl(imageUrl);
            if (StringUtils.isBlank(fileName)) {
                continue;
            }
            File dest = new File(baseDir, fileName);
            if (dest.exists()) {
                handled.incrementAndGet();
                System.out.println(imageUrl + "已存在");
                continue;
            }
            saveToFile(logFileName, job.toJSONString());
            // 异步处理
//            downloadFile(imageUrl);
            // 同步处理
            try {
                semaphore.acquire();
                download(imageUrl);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        System.out.println("已有文件数：" + handled.get());
        while (handled.get() < size) {
            try {
                Thread.sleep(1000L);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        System.out.println(handled.get());
        System.exit(1);
    }

    public static void saveToFile(String fileName, String txt) {
        try {
            FileWriter writer = new FileWriter(fileName, true); // 创建一个 FileWriter 对象来写入文件
            writer.write(txt);  // 写入文本内容
            writer.write("\r\n");
            writer.close();   // 关闭 FileWriter 对象
            System.out.println(fileName + "已成功保存。");
        } catch (IOException e) {
            System.out.println("文件保存时出错。");
            e.printStackTrace();
        }
    }

    private static String getFileNameFromUrl(String url) {
        try {
            return url.substring(27).replace("/", "_");
        } catch (Throwable e) {
            e.printStackTrace();
            return "";
        }
    }

    private static final OkHttpClient client = new OkHttpClient();
    private static final int REQUEST_EXTERNAL_STORAGE = 1;

    //下载文件
    private static void downloadFile(final String url) {
        new Thread(() -> download(url)).start();
    }

    //下载
    public static void download(final String url) {

        // System.out.println("DOWNLOAD" + "startTime=" + startTime);
        Request request = new Request.Builder().url(url)
                .addHeader("authority", "cdn.midjourney.com")
                .addHeader("accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7")
//                .addHeader("accept-language", "zh-CN,zh;q=0.9")
//                .addHeader("cache-control", "max-age=0")
//                .addHeader("sec-ch-ua", "\"Not.A/Brand\";v=\"8\", \"Chromium\";v=\"114\", \"Google Chrome\";v=\"114\"")
//                .addHeader("sec-ch-ua-mobile", "?0")
//                .addHeader("sec-ch-ua-platform", "\"Windows\"")
//                .addHeader("sec-fetch-dest", "document")
//                .addHeader("sec-fetch-mode", "navigate")
//                .addHeader("sec-fetch-site", "none")
//                .addHeader("sec-fetch-user", "?1")
//                .addHeader("upgrade-insecure-requests", "0")
//                .addHeader("user-agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36")
//                .addHeader("cookie", "_ga=GA1.1.862279448.1680306197; __cf_bm=qdAaL0Fu18C6rgdRFkCNbR8VHfwdeZfU6IsoLDUAwzE-1687962587-0-ASKyME62gVdsbRGCfijTqiqcNaXnSMgm51hHNiVSngqA4UdvUa3oslUnphqrnO6z4TKH7lF2EoQY2W43AEUCaYOY/cWmwIv4YnZpmn/oIh5; _ga_Q0DQ5L7K0D=GS1.1.1687962587.70.1.1687962676.0.0.0")
//                .addHeader("referer", "https://www.midjourney.com/showcase/recent/")
//                .addHeader("if-none-match", "\"6f9a7eabd54477316052d8ecac2ce170\"")
//                .addHeader("if-modified-since", "Tue, 27 Jun 2023 15:38:15 GMT")
                .build();
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                handled.incrementAndGet();
                // 下载失败
                e.printStackTrace();
                saveToFile(logFileName, "failed: " + url);
                System.out.println("download failed: " + url);
                semaphore.release();
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                if (response.code() != 200) {
                    System.out.println(url + " 失敗" + response.code());
                }
                Sink sink = null;
                BufferedSink bufferedSink = null;
                try {
                    String fileName = getFileNameFromUrl(url);
                    if (StringUtils.isBlank(fileName)) {
                        return;
                    }
                    File dest = new File(baseDir, fileName); // url.substring(27, url.lastIndexOf("/")) + ".png"
                    if (dest.exists()) {
                        System.out.println(url + "已存在");
                        return;
                    }
                    final long startTime = System.currentTimeMillis();
                    sink = Okio.sink(dest);
                    bufferedSink = Okio.buffer(sink);
                    bufferedSink.writeAll(response.body().source());

                    bufferedSink.close();
                    System.out.println(url + " download success " + "totalTime=" + (System.currentTimeMillis() - startTime));
                } catch (Exception e) {
                    e.printStackTrace();
                    System.out.println("download failed" + url);
                    saveToFile(logFileName, "failed: " + url);
                } finally {
                    handled.incrementAndGet();
                    if (bufferedSink != null) {
                        bufferedSink.close();
                    }
                    response.close();
                    semaphore.release();
                }
            }
        });
    }
}
